@claude review these changes, it it safe to mutate the AST in-place or will that cause issues? Is there a cleaner way to do it?

Claude finished @jansel's task in 7m 26s —— View job

Review: ellipsis indexing in normal mode

• Read the PR diff
• Analyze whether in-place AST mutation is safe
• Assess whether there's a cleaner approach

@jansel Here's my analysis of the two questions.

Is it safe to mutate the AST in place?

Yes, in this codebase it's safe — but for reasons worth making explicit, because the safety is load-bearing and not obvious from the diff.

1. The AST is per-compilation and freshly parsed. KernelCompiler.parse() calls inspect.getsource(fn) + ast.parse(source) on every compile() (kernel_compiler.py:90-94). Each BoundKernel gets its own tree, so mutating it can't leak across shape specializations or into the source cache (Kernel.source is @functools.cached on the text, not the tree). No shared/cached AST is being clobbered.

2. The mutation is intentionally persisted across passes — it's not just a local convenience. The device-IR lowering re-reads node.slice in _subscript_slice_proxy (device_ir.py:2625) and visit_Tuple. So the expanded slice must survive type propagation for codegen to see the right index list. This is actually the main argument for mutating the AST rather than expanding a throwaway key list: both passes read the one tree, so you write once.

3. It must be idempotent, and it is. For a store target the subscript is visited more than once — _assign calls self.visit(lhs) (→ visit_Subscript, expands) at type_propagation.py:502, then self.visit(lhs.slice) at :504. After the first expansion the slice is an ellipsis-free ast.Tuple, so the second pass is a no-op. Good, but this invariant is implicit — worth a one-line comment.

The real cost is altitude: visit_Subscript is part of a read-only type-inference NodeVisitor, and quietly rewriting the tree from inside it is a surprising side effect. It works, but the next reader won't expect a visitor to mutate.

Is there a cleaner way?

The natural instinct — "do it as a separate NodeTransformer pre-pass before type propagation" — doesn't work, because expansion needs the value tensor's ndim, which isn't known until the tensor type is resolved. So doing it during type propagation is a defensible place. Two cleaner-shaped options:

• Expand the resolved index list, not the AST. Handle ... in _device_indexing_size (type_info.py:341) / propagate_setitem and in compute_shape + _subscript_slice_proxy on the proxy list where ndim is known. Pro: type propagation stays read-only, no spooky tree mutation. Con: two-plus sites instead of one (loses the "both passes read one tree" benefit above).
• If keeping the AST mutation, make it explicit: rename to make the write obvious, add the idempotency comment, and factor the "expand ... given ndim" logic into a small pure helper so it's testable in isolation.

My recommendation: the AST-mutation approach is acceptable given the ndim constraint and the shared-tree read in device_ir. Just document the two invariants (fresh-per-compile + idempotent) at the call site.

Correctness issues in _expand_ellipsis_in_subscript

1. Negative n_expand silently drops the ellipsis. If the user over-indexes (dims_consumed > ndim), n_expand < 0, range(n_expand) is empty, and sl.elts[idx:idx+1] = [] deletes the ... with no error. torch would normally raise on over-indexing — here you'd get a confusing downstream failure instead. Consider asserting n_expand >= 0 and raising a clear error.

2. Multiple ellipses are handled inconsistently. Only ellipsis_indices[0] is expanded; a second ... stays as a raw ast.Constant in the tuple and hits codegen as a literal. Python's own rule is "an index can only have a single ellipsis" — better to detect len(ellipsis_indices) > 1 and raise that explicitly.

3. elt.value in (..., None) uses equality, not identity. Prefer elt.value is ... or elt.value is None. For an odd constant value, in triggers == which can misbehave (e.g. array-valued constants); is matches the singleton semantics you actually want and is consistent with line 768/780 which already use is ....

4. StackTensorType isn't covered. The guard only fires for TensorType (:761), but _assign treats StackTensorType as subscriptable too (:501). Ellipsis on a stack tensor won't expand. May be out of scope, but worth a note/test.

Test coverage

test_ellipsis_indexing covers leading [..., i] on a rank-3 tensor for both getter and setter — good. Not covered: trailing (x[i, ...]), middle, ellipsis-with-None, bare x[...], multiple ellipses (error path), and over-indexing (error path). The bare x[...] branch (:768) and the None-counting logic have no test exercising them.

Verdict

Approach is safe and reasonable. Before merge I'd suggest: (a) a comment documenting the fresh-per-compile + idempotency invariants, (b) explicit errors for multiple-ellipsis and over-indexing, (c) is instead of in (..., None), and (d) a couple more test cases for the trailing/None/error paths.
• branch fix-ellipsis-in-normal-mode